from typing import List, Dict, Any

import numpy as np
from langchain.prompts.example_selector.base import BaseExampleSelector


class CustomExampleSelector(BaseExampleSelector):

    def __init__(self, examples: List[Dict[str, str]]):
        self.examples = examples

    def add_example(self, example: Dict[str, str]) -> Any:
        """Add new example to store for a key."""
        self.examples.append(example)

    def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:
        """select which examples to use based on the inputs."""
        return np.random.choice(self.examples, size=2, replace=False)


if __name__ == '__main__':
    example_list = [
        {"foo": "1"},
        {"foo": "2"},
        {"foo": "3"}
    ]
    example_selector = CustomExampleSelector(example_list)
    example_selector.select_examples({"foo": "foo"})
    example_selector.add_example({"foo": "4"})
    print(example_selector.examples)